Helper Function

# Imports 
import numpy as np
import sympy as sym
from sympy.abc import t
%matplotlib inline
import matplotlib.pyplot as plt
import scipy.linalg
import time

# Helper 1

def hat(w,use_sym=True):
    if use_sym:
        what = sym.Matrix([[    0,-w[2], w[1]],
                           [ w[2],    0,-w[0]],
                           [-w[1], w[0],   0]])
    else:
        what = np.array([[    0,-w[2], w[1]],
                         [ w[2],    0,-w[0]],
                         [-w[1], w[0],   0]])
    return what

def unhat(what,use_sym=True): 
    if use_sym:
        w = sym.Matrix([what[0,3],what[1,3],what[2,3], what[2,1],what[0,2],what[1,0]])
    else:
        w = np.array([what[0,3],what[1,3],what[2,3], what[2,1],what[0,2],what[1,0]])
    return w

# def rot(w,theta,use_sym=True):
#     if use_sym:
#         rotMat = sym.Matrix(sym.simplify(sym.exp(hat(w,use_sym)*theta)))
#         for i in range(rotMat.shape[0]):
#             for j in range(rotMat.shape[1]):
#                 rotMat[i,j] = sym.simplify(rotMat[i,j].rewrite(sym.sin)) # Simplification made because sympy seems to have
#                                                                          # issues applying the Euler identity so I force it
#                                                                          # to recognize it.
#     else:
#         rotMat = scipy.linalg.expm(hat(w,use_sym)*theta)
#     return rotMat

def T(w,th,p,use_sym=True):
    R = rot(w,th,use_sym)
    if use_sym:
        Tmat = sym.Matrix([[R[0,0],R[0,1],R[0,2],p[0]],
                           [R[1,0],R[1,1],R[1,2],p[1]],
                           [R[2,0],R[2,1],R[2,2],p[2]],
                           [     0,     0,     0,   1]])
    else:
        Tmat = np.array([[R[0,0],R[0,1],R[0,2],p[0]],
                         [R[1,0],R[1,1],R[1,2],p[1]],
                         [R[2,0],R[2,1],R[2,2],p[2]],
                         [     0,     0,     0,   1]])
    return Tmat

def pos(T,use_sym=True):
    if use_sym:
        p = sym.Matrix([T[0,3],T[1,3],T[2,3]])
    else:
        p = np.array([T[0,3],T[1,3],T[2,3]])
    return p

def rot(th):
    rot_mat = sym.Matrix([[sym.cos(th),-sym.sin(th)],[sym.sin(th),sym.cos(th)]])
    return rot_mat

def rotNp(th):
    rot_mat = np.array([[np.cos(th),-np.sin(th)],[np.sin(th),np.cos(th)]])
    return rot_mat

#function to make sympy transformation matrix
def trans(R,p):
    g = sym.Matrix([[R[0],R[1],0,p[0]],[R[2],R[3],0,p[1]],[0,0,1,p[2]],[0,0,0,1]])
    return g

def transNp(R,p):
    g = np.array([[R[0],R[1],0,p[0]],[R[2],R[3],0,p[1]],[0,0,1,p[2]],[0,0,0,1]])
    return g

def trans_inverse(g):
    
  R_trans = g[:3,:3].T
  p = sym.Matrix([[g[0,3]], [g[1,3]], [g[2,3]]])
  RTp = -R_trans@p
  g_inv = sym.Matrix([[R_trans[0,0], R_trans[0,1], R_trans[0,2], RTp[0]], 
                  [R_trans[1,0], R_trans[1,1], R_trans[1,2], RTp[1]], 
                  [R_trans[2,0], R_trans[2,1], R_trans[2,2], RTp[2]],
                  [           0,            0,            0,     1]])

  return g_inv

import sympy as sym
from sympy import cos, sin, pi
import numpy as np

#Init of variables
t = sym.symbols('t')

M = 10   #Mass of Box
m = 1.5 #Mass of dice

L = 3   #Length of Box side 
l = 0.3   #Length of dice side

J_box = (1/6)*M*L**2
J_dice = (1/6)*m*l**2

g = 9.8 #Gravity
 
#Init of theta variables and their respective derivatives for box and dice

tta_box = sym.Function('theta_1')(t) #Theta box
tta_dice = sym.Function('theta_2')(t) #Theta dice


#Init of x/y variables and their respective derivatives for box
x_box = sym.Function('x_box')(t)         
y_box = sym.Function('y_box')(t)


#Init of x/y variables and their respective derivatives for dice
x_dice = sym.Function('x_dice')(t)         
y_dice = sym.Function('y_dice')(t)

#Array of DOF
q = sym.Matrix([x_box, y_box, tta_box, x_dice, y_dice, tta_dice])
qdot = q.diff(t)
qddot = qdot.diff(t)

Transform Phase

#Transforms for Box 
gwa_box = trans(rot(q[2]), sym.Matrix([q[0], q[1], 0])) #World to Centroid
gab_box = trans(rot(0),[L/2, 0, 0])                     #Shift to the right on x-axis
gac_box = trans(rot(0),[-L/2, 0, 0])                    #Shift to the left on the x-axis
gad_box = trans(rot(0),[0, L/2, 0])                     #Shift upwards on the y-axis
gae_box = trans(rot(0),[0, -L/2, 0])                    #Shift downwards on the y-axis


#Transforms for Dice 
gwa_dice = trans(rot(q[5]), sym.Matrix([q[3], q[4], 0])) #World to Centroid
gab_dice = trans(rot(0),[l/2, 0, 0])                     #Shift to the right on x-axis
gac_dice = trans(rot(0),[-l/2, 0, 0])                    #Shift to the left on the x-axis
gad_dice = trans(rot(0),[0, l/2, 0])                     #Shift upwards on the y-axis
gae_dice = trans(rot(0),[0, -l/2, 0])                    #Shift downwards on the y-axis

#Transforms World to Object

#World to Box
gwb_box = gwa_box@gab_box       #World to right of box
gwc_box = gwa_box@gac_box       #World to left of box
gwd_box = gwa_box@gad_box       #World to up of box
gwe_box = gwa_box@gae_box       #World to down of box

#World to Dice
gwb_dice = gwa_dice@gab_dice    #World to right of box
gwc_dice = gwa_dice@gac_dice    #World to left of box
gwd_dice = gwa_dice@gad_dice    #World to up of box
gwe_dice = gwa_dice@gae_dice    #World to down of box


#Box to Dice 
gaa_boxDice = trans_inverse(gwa_box)@gwa_dice           #Box center to Dice
gba_boxDice = trans_inverse(gwb_box)@gwa_dice           #Box right to Dice
gca_boxDice = trans_inverse(gwc_box)@gwa_dice           #Box left to Dice
gda_boxDice = trans_inverse(gwd_box)@gwa_dice           #Box up to Dice
gea_boxDice = trans_inverse(gwe_box)@gwa_dice           #Box down to Dice

#Points on Dice
pDiceRight = sym.Matrix([l/2, 0, 0, 1])
pDiceLeft = sym.Matrix([-l/2, 0, 0, 1])
pDiceUp = sym.Matrix([0,l/2, 0, 1])
pDiceDown = sym.Matrix([0,-l/2, 0, 1])

#Points on Dice relative to Box
pBoxDiceRight1 = gba_boxDice@pDiceRight
pBoxDiceLeft1 = gba_boxDice@pDiceLeft
pBoxDiceUp1 = gba_boxDice@pDiceUp
pBoxDiceDown1 = gba_boxDice@pDiceDown

pBoxDiceRight2 = gca_boxDice@pDiceRight
pBoxDiceLeft2 = gca_boxDice@pDiceLeft
pBoxDiceUp2 = gca_boxDice@pDiceUp
pBoxDiceDown2 = gca_boxDice@pDiceDown


pBoxDiceRight3 = gda_boxDice@pDiceRight
pBoxDiceLeft3 = gda_boxDice@pDiceLeft
pBoxDiceUp3 = gda_boxDice@pDiceUp
pBoxDiceDown3 = gda_boxDice@pDiceDown


pBoxDiceRight4 = gea_boxDice@pDiceRight
pBoxDiceLeft4 = gea_boxDice@pDiceLeft
pBoxDiceUp4 = gea_boxDice@pDiceUp
pBoxDiceDown4 = gea_boxDice@pDiceDown
#Phi impacts
# phiRx_plus = pBoxDiceRight[0] + L/2
# phiRx_minus = pBoxDiceRight[0] - L/2
# phiRy_plus = pBoxDiceRight[1] + L/2
# phiRy_minus = pBoxDiceRight[1] - L/2

# phiLx_plus = pBoxDiceLeft[0] + L/2
# phiLx_minus = pBoxDiceLeft[0] - L/2
# phiLy_plus = pBoxDiceLeft[1] + L/2
# phiLy_minus = pBoxDiceLeft[1] - L/2

# phiUx_plus = pBoxDiceUp[0] + L/2
# phiUx_minus = pBoxDiceUp[0] - L/2
# phiUy_plus = pBoxDiceUp[1] + L/2
# phiUy_minus = pBoxDiceUp[1] - L/2

# phiDx_plus = pBoxDiceDown[0] + L/2
# phiDx_minus = pBoxDiceDown[0] - L/2
# phiDy_plus = pBoxDiceDown[1] + L/2
# phiDy_minus = pBoxDiceDown[1] - L/2

phiRx_plus = pBoxDiceRight1[0] 
phiRx_minus = pBoxDiceRight1[0]
phiRy_plus = pBoxDiceRight1[0]
phiRy_minus = pBoxDiceRight1[0]

phiLx_plus = pBoxDiceLeft2[0]
phiLx_minus = pBoxDiceLeft2[0]
phiLy_plus = pBoxDiceLeft2[0]
phiLy_minus = pBoxDiceLeft2[0]

phiUx_plus = pBoxDiceUp3[1]
phiUx_minus = pBoxDiceUp3[1]
phiUy_plus = pBoxDiceUp3[1]
phiUy_minus = pBoxDiceUp3[1]

phiDx_plus = pBoxDiceDown4[1]
phiDx_minus = pBoxDiceDown4[1]
phiDy_plus = pBoxDiceDown4[1]
phiDy_minus = pBoxDiceDown4[1]

Phi Impact Variables

phi_list = [phiRx_plus, phiRx_minus, phiRy_plus, phiRy_minus, 
            phiLx_plus, phiLx_minus, phiLy_plus, phiLy_minus, 
            phiUx_plus, phiUx_minus, phiUy_plus, phiUy_minus, 
            phiDx_plus, phiDx_minus, phiDy_plus, phiDy_minus]

phi_fun_1 = sym.lambdify([*q,*qdot], phiRx_plus, modules=sym)
phi_fun_2 = sym.lambdify([*q,*qdot], phiRx_minus, modules=sym)
phi_fun_3 = sym.lambdify([*q,*qdot], phiRy_plus, modules=sym)
phi_fun_4 = sym.lambdify([*q,*qdot], phiRy_minus, modules=sym)

phi_fun_5 = sym.lambdify([*q,*qdot], phiLx_plus, modules=sym)
phi_fun_6 = sym.lambdify([*q,*qdot], phiLx_minus, modules=sym)
phi_fun_7 = sym.lambdify([*q,*qdot], phiLy_plus, modules=sym)
phi_fun_8 = sym.lambdify([*q,*qdot], phiLy_minus, modules=sym)

phi_fun_9 = sym.lambdify([*q,*qdot], phiUx_plus, modules=sym)
phi_fun_10 = sym.lambdify([*q,*qdot], phiUx_minus, modules=sym)
phi_fun_11 = sym.lambdify([*q,*qdot], phiUy_plus, modules=sym)
phi_fun_12 = sym.lambdify([*q,*qdot], phiUx_minus, modules=sym)

phi_fun_13 = sym.lambdify([*q,*qdot], phiDx_minus, modules=sym)
phi_fun_14 = sym.lambdify([*q,*qdot], phiDy_plus, modules=sym)
phi_fun_15 = sym.lambdify([*q,*qdot], phiDx_minus, modules=sym)
phi_fun_16 = sym.lambdify([*q,*qdot], phiDy_plus, modules=sym)

phi_func_list = [phi_fun_1, phi_fun_2, phi_fun_3, phi_fun_4, 
                 phi_fun_5, phi_fun_6, phi_fun_7, phi_fun_8,
                 phi_fun_9, phi_fun_10, phi_fun_11, phi_fun_12,
                 phi_fun_13, phi_fun_14, phi_fun_15, phi_fun_16]

Impact Conditions

gwa_inverse_box = trans_inverse(gwa_box)        #World to box inverse
gwa_inverse_dice = trans_inverse(gwa_dice)      #World to dice inverse

va_hat_box = gwa_inverse_box@gwa_box.diff(t)     #hat of box centroid
va_hat_dice = gwa_inverse_dice@gwa_dice.diff(t)  #hat of dice centroid

Vb_unhat_box = unhat(va_hat_box, use_sym = True) #Velocity box
Vb_unhat_dice= unhat(va_hat_dice, use_sym = True) #Velocity dice

display(Vb_unhat_box)
display(Vb_unhat_dice)

mI_box = np.array([[M, 0, 0, 0, 0, 0],
                   [0, M, 0, 0, 0, 0],
                   [0, 0, M, 0, 0, 0],
                   [0, 0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0, J_box]])             #mass matrix

mI_dice = np.array([[m, 0, 0, 0, 0, 0],
                   [0, m, 0, 0, 0, 0],
                   [0, 0, m, 0, 0, 0],
                   [0, 0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0, 0],
                   [0, 0, 0, 0, 0, J_dice]])

$\displaystyle \left[\begin{matrix}\sin{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{y_{box}}{\left(t \right)} + \cos{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{x_{box}}{\left(t \right)}\\- \sin{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{x_{box}}{\left(t \right)} + \cos{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{y_{box}}{\left(t \right)}\\0\\0\\0\\\sin^{2}{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \theta_{1}{\left(t \right)} + \cos^{2}{\left(\theta_{1}{\left(t \right)} \right)} \frac{d}{d t} \theta_{1}{\left(t \right)}\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\sin{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{y_{dice}}{\left(t \right)} + \cos{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{x_{dice}}{\left(t \right)}\\- \sin{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{x_{dice}}{\left(t \right)} + \cos{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \operatorname{y_{dice}}{\left(t \right)}\\0\\0\\0\\\sin^{2}{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \theta_{2}{\left(t \right)} + \cos^{2}{\left(\theta_{2}{\left(t \right)} \right)} \frac{d}{d t} \theta_{2}{\left(t \right)}\end{matrix}\right]$

Kinetic Energy, Pottential, Lagrangian

Ke_box = (0.5*Vb_unhat_box.T*mI_box*Vb_unhat_box)[0]
Ke_dice = (0.5*Vb_unhat_dice.T*mI_dice*Vb_unhat_dice)[0]
Ke_tot = Ke_box+Ke_dice

Pe_box =  -M*g*gwa_box[1,3]
Pe_dice = -m*g*gwa_dice[1,3]
Pe_tot = Pe_box+Pe_dice

L = Ke_tot+Pe_tot
L_mat = sym.Matrix([L])
display(sym.simplify(L_mat))

$\displaystyle \left[\begin{matrix}- 98.0 \operatorname{y_{box}}{\left(t \right)} - 14.7 \operatorname{y_{dice}}{\left(t \right)} + 7.5 \left(\frac{d}{d t} \theta_{1}{\left(t \right)}\right)^{2} + 0.01125 \left(\frac{d}{d t} \theta_{2}{\left(t \right)}\right)^{2} + 5.0 \left(\frac{d}{d t} \operatorname{x_{box}}{\left(t \right)}\right)^{2} + 0.75 \left(\frac{d}{d t} \operatorname{x_{dice}}{\left(t \right)}\right)^{2} + 5.0 \left(\frac{d}{d t} \operatorname{y_{box}}{\left(t \right)}\right)^{2} + 0.75 \left(\frac{d}{d t} \operatorname{y_{dice}}{\left(t \right)}\right)^{2}\end{matrix}\right]$

Phi Constraints

Euler-Lagrange

dldq = L_mat.jacobian(q)
dldqdot = L_mat.jacobian(qdot)
dldqdot_dt = dldqdot.diff(t) 
lhs = sym.simplify(dldqdot_dt.T-dldq.T)
rhs = sym.Matrix([0, 108, 0, 0, 0, 0])
EL = sym.Eq(lhs,rhs)
display(EL)
EL_sol = sym.solve(EL,qddot)
display(EL_sol)

$\displaystyle \left[\begin{matrix}10.0 \frac{d^{2}}{d t^{2}} \operatorname{x_{box}}{\left(t \right)}\\10.0 \frac{d^{2}}{d t^{2}} \operatorname{y_{box}}{\left(t \right)} + 98.0\\15.0 \frac{d^{2}}{d t^{2}} \theta_{1}{\left(t \right)}\\1.5 \frac{d^{2}}{d t^{2}} \operatorname{x_{dice}}{\left(t \right)}\\1.5 \frac{d^{2}}{d t^{2}} \operatorname{y_{dice}}{\left(t \right)} + 14.7\\0.0225 \frac{d^{2}}{d t^{2}} \theta_{2}{\left(t \right)}\end{matrix}\right] = \left[\begin{matrix}0\\108\\0\\0\\0\\0\end{matrix}\right]$

{Derivative(x_box(t), (t, 2)): 0.0,
 Derivative(y_box(t), (t, 2)): 1.00000000000000,
 Derivative(theta_1(t), (t, 2)): 0.0,
 Derivative(x_dice(t), (t, 2)): 0.0,
 Derivative(y_dice(t), (t, 2)): -9.80000000000000,
 Derivative(theta_2(t), (t, 2)): 0.0}

Lambdify

xddot_box_fun = sym.lambdify([*q, *qdot], EL_sol[qddot[0]], modules = sym)

yddot_box_fun = sym.lambdify([*q, *qdot], EL_sol[qddot[1]], modules = sym)

ttaddot_box_fun = sym.lambdify([*q, *qdot],EL_sol[qddot[2]], modules = sym)

xddot_dice_fun = sym.lambdify([*q, *qdot], EL_sol[qddot[3]], modules = sym)

yddot_dice_fun = sym.lambdify([*q, *qdot], EL_sol[qddot[4]], modules = sym)

ttaddot_dice_fun = sym.lambdify([*q, *qdot], EL_sol[qddot[5]], modules = sym)

Before and After Impact

xmin_dice, ymin_dice, ttamin_dice = sym.symbols('x_diceMin, y_diceMin, theta_diceMin')
xmin_box, ymin_box, ttamin_box = sym.symbols('x_boxMin, y_boxMin, theta_boxMin')

xdotplus_box, ydotplus_box, ttadotplus_box, xdotmin_box, ydotmin_box, ttadotmin_box = sym.symbols('xdot_boxPlus, ydot_boxPlus, thetadot_boxPlus, xdot_boxMin, ydot_boxMin, thetadot_boxMin')
xdotplus_dice, ydotplus_dice, ttadotplus_dice, xdotmin_dice, ydotmin_dice, ttadotmin_dice = sym.symbols('xdot_dicePlus, ydot_dicePlus, thetadot_dicePlus, xdot_diceMin, ydot_diceMin, thetadot_diceMin')
lamb = sym.symbols('lambda')


display(q[0])
tau_minus = {q[0]:xmin_box, q[1]:ymin_box, q[2]:ttamin_box, q[3]:xmin_dice, q[4]:ymin_dice, q[5]:ttamin_dice, 
            qdot[0]:xdotmin_box, qdot[1]:ydotmin_box, qdot[2]:ttadotmin_box, qdot[3]:xdotmin_dice, qdot[4]:ydotmin_dice, qdot[5]:ttadotmin_dice}

display(q[0])

tau_plus = {q[0]:xmin_box, q[1]:ymin_box, q[2]:ttamin_box, q[3]:xmin_dice, q[4]:ymin_dice, q[5]:ttamin_dice,
            qdot[0]:xdotplus_box, qdot[1]:ydotplus_box, qdot[2]:ttadotplus_box, qdot[3]:xdotplus_dice, qdot[4]:ydotplus_dice, qdot[5]:ttadotplus_dice}


dldqdot_plus = dldqdot.subs(tau_plus)
dldqdot_minus = dldqdot.subs(tau_minus)

H = dldqdot*qdot-L_mat
H_minus = H.subs(tau_minus)
H_plus = H.subs(tau_plus)

delta_dldqdot = dldqdot_plus-dldqdot_minus
delta_dldqdot = delta_dldqdot.T
delta_dldqdot = sym.simplify(delta_dldqdot)


delta_H= H_plus-H_minus
delta_H= sym.simplify(delta_H)

$\displaystyle \operatorname{x_{box}}{\left(t \right)}$

$\displaystyle \operatorname{x_{box}}{\left(t \right)}$

Impact Update

def impact_update(s, phi):
    impact_lhs = sym.Matrix([delta_dldqdot, delta_H[0]])
    dphidq = phi.diff(q)
    impact_rhs = sym.Matrix([lamb*dphidq, 0])
    impactEq = sym.Eq(impact_lhs, impact_rhs)
    numerical = {q[0]: s[0],
                 q[1]: s[1],
                 q[2]: s[2],
                 q[3]: s[3],
                 q[4]: s[4],
                 q[5]: s[5],
                 xdotmin_box: s[6],
                 ydotmin_box: s[7],
                 ttadotmin_box: s[8],
                 xdotmin_dice: s[9],
                 ydotmin_dice: s[10],
                 ttadotmin_dice: s[11]}

    impactEq = impactEq.subs(numerical)
    print('\nImpact\n')
    solvedState = sym.solve(impactEq,[xdotplus_box, ydotplus_box, ttadotplus_box, 
                                      xdotplus_dice,ydotplus_dice,ttadotplus_dice, lamb], dict=True)

    if len(solvedState) > 1:
        for i in solvedState:
            if abs(i[lamb]) < 1e-06:
                pass
            else:
                return np.array([s[0],s[1],s[2],s[3],s[4],s[5],
                                i[xdotplus_box],
                                i[ydotplus_box],
                                i[ttadotplus_box],
                                i[xdotplus_dice],
                                i[ydotplus_dice],
                                i[ttadotplus_dice]])


# dpdq_minus = dpdq.subs(subs_minus)

def impact_condition(x, thresh = 0.1):
    for i in range(len(phi_func_list)):
        phi_val = phi_func_list[i]
        if phi_val(*x) > -thresh and phi_val(*x) < thresh:
            print('Impact detected')
            return phi_list[i]
    return False


def simulate_with_impact(f, x0, tspan, dt, integrate):
    
    N = int((max(tspan)-min(tspan))/dt)
    x = np.copy(x0)
    tvec = np.linspace(min(tspan),max(tspan),N)
    xtraj = np.zeros((len(x0),N))

    for i in range(N):
        phi = impact_condition(x)        
        if  phi != False:
            x = impact_update(x, phi)
            xtraj[:,i]=integrate(f,x,dt)
        else:
            xtraj[:,i]=integrate(f,x,dt)
        x = np.copy(xtraj[:,i])
    return xtraj 

def integrate(f,x0,dt):

    k1=dt*f(x0) 
    k2=dt*f(x0+k1/2.)
    k3=dt*f(x0+k2/2.)
    k4=dt*f(x0+k3)
    xnew=x0+(1/6.)*(k1+2.0*k2+2.0*k3+k4)
    return xnew 
    

def dyn(s):

    sdot1 = s[6]
    sdot2 = s[7]
    sdot3 = s[8]
    sdot4 = s[9]
    sdot5 = s[10]
    sdot6 = s[11]
    sdot7 = xddot_box_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])
    sdot8 = yddot_box_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])
    sdot9 = ttaddot_box_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])
    sdot10 = xddot_dice_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])
    sdot11 = yddot_dice_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])
    sdot12 = ttaddot_dice_fun(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10], s[11])

    return np.array([sdot1, sdot2, sdot3, sdot4, sdot5, sdot6, sdot7, sdot8, sdot9, sdot10, sdot11, sdot12])

s0 = [0, 0, 0, 0, 0, np.pi/4, 0, 0, 0, 0, 0, 0]

traj = simulate_with_impact(dyn, s0, [0,10], 0.01, integrate)

x = np.arange(0,10,0.01)

plt.plot(x, traj[0].T, label = 'x box')
plt.plot(x, traj[1].T, label = 'y box')
plt.plot(x, traj[2].T, label = 'theta box')
plt.plot(x, traj[3].T, label = 'x dice')
plt.plot(x, traj[4].T, label = 'y dice')
plt.plot(x, traj[5].T, label = 'theta dice')

plt.legend()
plt.show()

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Impact detected

Impact

Animate Function

def animate_BoxandDice(array, L=3,l=0.3,T=3):
    """
    Function to generate web-based animation of dice in a box

    Parameters:
    ================================================
    theta_array:
        trajectory of theta1 and theta2, should be a NumPy array with
        shape of (2,N)
    L1:
        length of the first pendulum
    L2:
        length of the second pendulum
    T:
        length/seconds of animation duration

    Returns: None
    """

    ################################
    # Imports required for animation.
    from plotly.offline import init_notebook_mode, iplot
    from IPython.display import display, HTML
    import plotly.graph_objects as go

    #######################
    # Browser configuration.
    def configure_plotly_browser_state():
        import IPython
        display(IPython.core.display.HTML('''
            <script src="/static/components/requirejs/require.js"></script>
            <script>
              requirejs.config({
                paths: {
                  base: '/static/base',
                  plotly: 'https://cdn.plot.ly/plotly-1.5.1.min.js?noext',
                },
              });
            </script>
            '''))
    configure_plotly_browser_state()
    init_notebook_mode(connected=False)

    ###############################################
    # Define arrays containing data for frame axes
    # In each frame, the x and y axis are always fixed

    x_axis = np.array([0.3, 0.0])
    y_axis = np.array([0.0, 0.3])


    N = len(array[0]) # Need this for specifying length of simulation

    ###############################################
    # Define arrays containing data for frame axes
    # In each frame, the x and y axis are always fixed
    x_axis = np.array([0.3, 0.0])
    y_axis = np.array([0.0, 0.3])
    # Use homogeneous tranformation to transfer these two axes/points
    # back to the fixed frame

    ###### B Frame
    frame_b_x_axis = np.zeros((2,N))
    frame_b_y_axis = np.zeros((2,N))

    ###### C Frame
    frame_c_x_axis = np.zeros((2,N))
    frame_c_y_axis = np.zeros((2,N))

    # Frame Box Upper Right 
    frame_box_x_UR = np.zeros((1,N))
    frame_box_y_UR = np.zeros((1,N))

    # Frame Box Upper Left 
    frame_box_x_UL = np.zeros((1,N))
    frame_box_y_UL = np.zeros((1,N))

    # Frame Box Lower Right 
    frame_box_x_LR = np.zeros((1,N))
    frame_box_y_LR = np.zeros((1,N))

    # Frame Box Lower Left 
    frame_box_x_LL = np.zeros((1,N))
    frame_box_y_LL = np.zeros((1,N))

    # Frame c Upper Right 
    frame_dice_x_UR = np.zeros((1,N))
    frame_dice_y_UR = np.zeros((1,N))

    # Frame c Upper Left 
    frame_dice_x_UL = np.zeros((1,N))
    frame_dice_y_UL = np.zeros((1,N))

    # Frame c Lower Right 
    frame_dice_x_LR = np.zeros((1,N))
    frame_dice_y_LR = np.zeros((1,N))

    # Frame c Lower Left 
    frame_dice_x_LL = np.zeros((1,N))
    frame_dice_y_LL = np.zeros((1,N)) 
    

    for i in range(N): # iteration through each time step
        # evaluate homogeneous transformation

#Box 2d Construction

          #World to Centroid
          t_wa_box = np.array([[np.cos(array[2][i]),-np.sin(array[2][i]), 0, array[0][i]],
                               [np.sin(array[2][i]),np.cos(array[2][i]), 0, array[1][i]],
                               [0, 0, 1, array[2][i]],
                               [0, 0, 0, 1]])

          #Box Corners
          Pur_box = np.array([L/2, L/2, 0, 1])
          Pul_box = np.array([-L/2, L/2, 0, 1])
          Plr_box = np.array([L/2, -L/2, 0, 1])
          Pll_box = np.array([-L/2, -L/2, 0, 1])

        # Frame B Upper Right 
          box_UR = t_wa_box@Pur_box
          frame_box_x_UR[0,i] = box_UR[0]
          frame_box_y_UR[0,i] = box_UR[1]

        # Frame B Upper Left 
          box_UL = t_wa_box@Pul_box
          frame_box_x_UL[0,i] = box_UL[0]
          frame_box_y_UL[0,i] = box_UL[1]

        # Frame B Lower Right
          box_LR = t_wa_box@Plr_box
          frame_box_x_LR[0,i] = box_LR[0]
          frame_box_y_LR[0,i] = box_LR[1]

        # Frame B Lower Left 
          box_LL = t_wa_box@Pll_box
          frame_box_x_LL[0,i] = box_LL[0]
          frame_box_y_LL[0,i] = box_LL[1]

#Dice 2d Construction

          # t_wa_dice = transNp(rotNp(array[5][i]), np.array([array[3][i], array[4][i], array[5][i], 0])) #World to Centroid

          t_wa_dice = np.array([[np.cos(array[5][i]),-np.sin(array[5][i]), 0, array[3][i]],
                               [np.sin(array[5][i]),np.cos(array[5][i]), 0, array[4][i]],
                               [0, 0, 1, array[5][i]],
                               [0, 0, 0, 1]])

          Pup_dice = np.array([0, l/2, 0, 1])
          Pright_dice = np.array([l/2, 0, 0, 1])
          Pleft_dice = np.array([-l/2, 0, 0, 1])
          Pdown_dice = np.array([0, -l/2, 0, 1])

        # Frame c Upper Right 
          cUR = t_wa_dice@Pup_dice
          frame_dice_x_UR[0,i] = cUR[0]
          frame_dice_y_UR[0,i] = cUR[1]

        # Frame c Upper Left
          cUL = t_wa_dice@Pright_dice
          frame_dice_x_UL[0,i] = cUL[0]
          frame_dice_y_UL[0,i] = cUL[1]

        # Frame c Lower Right
          cLR = t_wa_dice@Pleft_dice
          frame_dice_x_LR[0,i] = cLR[0]
          frame_dice_y_LR[0,i] = cLR[1]

        # Frame c Lower Left
          cLL = t_wa_dice@Pdown_dice
          frame_dice_x_LL[0,i] = cLL[0]
          frame_dice_y_LL[0,i] = cLL[1]     
                
    ###################################
    # Using these to specify axis limits.
    
    xm = -3
    xM = 3
    ym = -3
    yM = 3

    ###########################
    # Defining data dictionary.
    # Trajectories are here.
    data=[
        # note that except for the trajectory (which you don't need this time),
        # you don't need to define entries other than "name". The items defined
        # in this list will be related to the items defined in the "frames" list
        # later in the same order. Therefore, these entries can be considered as 
        # labels for the components in each animation frame
        dict(name='Rectangle B'),
        dict(name='Rectangle C'),
        dict(name='Ground'),
        #######

        # You don't need to show trajectory this time,
        # but if you want to show the whole trajectory in the animation (like what
        # you did in previous homeworks), you will need to define entries other than 
        # "name", such as "x", "y". and "mode".

        # dict(x=xx1, y=yy1, 
        #      mode='markers', name='Pendulum 1 Traj', 
        #      marker=dict(color="fuchsia", size=2)
        #     ),
        # dict(x=xx2, y=yy2, 
        #      mode='markers', name='Pendulum 2 Traj', 
        #      marker=dict(color="purple", size=2)
        #     ),
        ]

    ################################
    # Preparing simulation layout.
    # Title and axis ranges are here.
    layout=dict(autosize=False, width=1000, height=1000,
                xaxis=dict(range=[xm, xM], autorange=False, zeroline=False,dtick=1),
                yaxis=dict(range=[ym, yM], autorange=False, zeroline=False,scaleanchor = "x",dtick=1),
                title='Double Pendulum Simulation', 
                hovermode='closest',
                updatemenus= [{'type': 'buttons',
                               'buttons': [{'label': 'Play','method': 'animate',
                                            'args': [None, {'frame': {'duration': T, 'redraw': False}}]},
                                           {'args': [[None], {'frame': {'duration': T, 'redraw': False}, 'mode': 'immediate',
                                            'transition': {'duration': 0}}],'label': 'Pause','method': 'animate'}
                                          ]
                              }]
               )

    ########################################
    # Defining the frames of the simulation.
    # This is what draws the lines from
    # joint to joint of the pendulum.
    frames=[dict(data=[# first three objects correspond to the arms and two masses,
                       # same order as in the "data" variable defined above (thus 
                       # they will be labeled in the same order)
                       dict(x=[frame_box_x_UR[0,k], frame_box_x_UL[0,k], frame_box_x_LL[0,k], frame_box_x_LR[0,k], frame_box_x_UR[0,k]], 
                            y=[frame_box_y_UR[0,k], frame_box_y_UL[0,k], frame_box_y_LL[0,k], frame_box_y_LR[0,k], frame_box_y_UR[0,k]], 
                            mode='lines',
                            line=dict(color='red', width=3),
                            ),

                       dict(x=[frame_dice_x_UR[0,k], frame_dice_x_UL[0,k], frame_dice_x_LL[0,k], frame_dice_x_LR[0,k], frame_dice_x_UR[0,k]], 
                            y=[frame_dice_y_UR[0,k], frame_dice_y_UL[0,k], frame_dice_y_LL[0,k], frame_dice_y_LR[0,k], frame_dice_y_UR[0,k]], 
                            mode='lines',
                            line=dict(color='blue', width=3),
                            ),

                       dict(x=[-3,3], 
                            y=[0,0], 
                            mode='lines',
                            line=dict(color='green', width=3),
                            ),                        
                      ]) for k in range(N)]

    #######################################
    # Putting it all together and plotting.
    figure1=dict(data=data, layout=layout, frames=frames)           
    iplot(figure1)

animate_BoxandDice(traj,L=3, l=0.3,T=3)